package pcap

import (
	"errors"
	"fmt"

	"github.com/kin9-0rz/gopcap"
)

// DNSInfo DNS信息
type DNSInfo struct {
	Domain  string      `json:"domain"`
	Answers []DNSAnswer `json:"answers"`
}

// NewDNSInfo new dns info
func NewDNSInfo(domain string) DNSInfo {
	return DNSInfo{
		Domain:  domain,
		Answers: make([]DNSAnswer, 0),
	}
}

// AddDNSAnswer add dns answer
func (d *DNSInfo) AddDNSAnswer(name string, addr string) {
	a := DNSAnswer{
		Name: name,
		Addr: addr,
	}
	d.Answers = append([]DNSAnswer{a}, d.Answers...)
}

// DNSAnswer DNS应答内容
type DNSAnswer struct {
	Name string `json:"name"`
	Addr string `json:"addr"`
}

// DNS 协议结构
type DNS struct {
	raw          []byte
	id           int
	flags        int
	countQueries int
	countAnswers int
	countAuthRR  int
	countAddRR   int
	queries      []string
	dnsInfo      DNSInfo
}

// FromBytes 解析DNS应答数据
func (d *DNS) FromBytes(data []byte) error {
	d.raw = data
	size := len(data)
	if size < 35 {
		return nil
	}
	// 解析头
	d.id = int(gopcap.GetUint16(data[0:2], false))
	d.flags = int(gopcap.GetUint16(data[2:4], false))

	d.countQueries = int(gopcap.GetUint16(data[4:6], false))
	d.countAnswers = int(gopcap.GetUint16(data[6:8], false))
	d.countAuthRR = int(gopcap.GetUint16(data[8:10], false))
	d.countAddRR = int(gopcap.GetUint16(data[10:12], false))

	if d.countQueries != 1 {
		return nil
	}

	// 解析查询内容，数量取决于countQueries
	d.queries = make([]string, d.countQueries)
	pos := 12

	if data[pos] == 0 {
		// 不处理<ROOT>这种类型
		return nil
	}

	// TODO 多个查询会导致异常，但是，找不到合适的pcap验证。
	for index := 0; index < d.countQueries; index++ {
		domain, off := d.readDomain(data[pos:])
		d.queries[index] = domain
		d.dnsInfo = NewDNSInfo(domain)
		pos += off

		// 查询类型，通常是A, 2
		// 查询类，通常是IN，2
		// 用不上，直接跳过
		pos += 4
	}

	// 应答成功，没错误
	if d.flags != 0x8180 {
		return nil
	}

	if d.countAnswers == 0 {
		return nil
	}

	if len(data) <= pos {
		return nil
	}

	// 解析应答内容，数量取决于countAnswers
	for index := 0; index < d.countAnswers; index++ {
		domain := ""
		offset := 0

		if data[pos]>>4<<4 == 0xc0 {
			// c 开头的话，则是直接索引
			offset := (int(data[pos]^0xc0))<<8 + int(data[pos+1])
			domain, _ = d.readDomain(data[offset:])
			pos += 2 // skip c0 ?? - name: 索引

			if data[pos] == 0 && data[pos+1] == 0x01 {
				// 类型A A(Host Address)
				// skip 00 10 - type:
				// skip 00 01 - class: in
				// skip 00 00 00 78 - time to live: 120
				pos += 8
				dataLength := int(gopcap.GetUint16(data[pos:pos+2], false))
				pos += 2
				addr := ConvertIP(data[pos : pos+dataLength])
				d.dnsInfo.AddDNSAnswer(domain, addr)
				pos += dataLength
				continue
			} else if data[pos] == 0 && data[pos+1] == 0x10 {
				// 类型 TXT
				pos += 8
				dataLength := int(gopcap.GetUint16(data[pos:pos+2], false))
				pos += 2 // skip data length
				d.dnsInfo.AddDNSAnswer(domain, string(data[pos+1:pos+dataLength]))
				pos += dataLength
				continue
			} else if data[pos] == 0 && data[pos+1] == 0x1c {
				// 类型 AAAA IPV6
				pos += 8
				dataLength := int(gopcap.GetUint16(data[pos:pos+2], false))
				pos += 2
				ipv6 := ""
				for i := 0; i < dataLength; i++ {
					ipv6 += fmt.Sprintf("%02x", data[pos+i])
					if i%2 == 1 && i != dataLength-1 {
						ipv6 += ":"
					}
				}
				d.dnsInfo.AddDNSAnswer(domain, ipv6)
				pos += dataLength
				continue
			} else if data[pos] == 0 && data[pos+1] == 0x05 {
				// CNAME(Canonical NAME for an alias)
				pos += 8
				dataLength := int(gopcap.GetUint16(data[pos:pos+2], false))
				pos += 2
				addrData := data[pos : pos+dataLength]
				addr, _ := d.readDomain(addrData)
				d.dnsInfo.AddDNSAnswer(domain, addr)
				pos += dataLength
				continue
			} else {
				return errors.New("Unknown DNS Type")
			}
		} else if data[pos] == 0 {
			if data[pos+1] == 0 && data[pos+2] == 0x6 {
				// SOA (Start of a zone of Authority)
				pos += 9
				dataLength := int(gopcap.GetUint16(data[pos:pos+2], false))
				pos += dataLength
				continue
			} else {
				return errors.New("Unknown DNS Type-<ROOT>")
			}
		} else {
			// 处理非索引数据
			domain, offset = d.readDomain(data[pos:])
			pos += offset // skip domain
		}
		pos += 8

		// 数据长度，解析的长度必然不会超过这个长度
		dataLength := int(gopcap.GetUint16(data[pos:pos+2], false))
		pos += 2

		addrData := data[pos : pos+dataLength]
		// 如果长度为4，则是IPv4，其他的则是cname地址
		addr := ""
		if dataLength == 4 {
			addr = ConvertIP(data[pos : pos+dataLength])
		} else {
			addr, _ = d.readDomain(addrData)
		}

		d.dnsInfo.AddDNSAnswer(domain, addr)

		pos += dataLength
	}

	return nil
}

// readDomain 读取一段可变长数据（域名），返回内容和偏移长度
func (d *DNS) readDomain(bytes []byte) (string, int) {
	domain := ""
	index := 0
	for {
		off := int(bytes[index])
		if off == 0 {
			domain = domain[:len(domain)-1]
			index++
			break
		} else if (off >> 4 << 4) == 0xc0 {
			index++
			offset := (int(off^0xc0))<<8 + int(bytes[index])
			end, _ := d.readDomain(d.raw[offset:])
			domain += end
			index++
			break
		}
		index++
		domain += string(bytes[index:index+off]) + "."
		index = index + off
	}

	return domain, index
}
